/// The ray tracer base function

#define TMAX 200.
#define TMIN.1
#define INF 1e9
#define PI 3.141592654
#define EPS 1e-9

#define WS_ // World Space
#define OS_ // Object Space
#define VS_ // Camera Space

#define P_  // point
#define V_  // vector 
#define N_  // normal

#define SCENE_// Scene Object

#define GI 0 // whether enable global illumination

const vec3 V_ L_S=vec3(1.,0.,0.);
const vec3 V_ L_U=vec3(0.,1.,0.);
const vec3 V_ L_F=vec3(0.,0.,1.);

///
/// Objects
///

// Simple plane
struct Plane{
    vec3 normal;
    float offset;
    uint matIndex;
};

struct Sphere{
    float radius;
    vec3 center;
    uint matIndex;
};

struct Rectangle{
    vec3 origin;
    vec3 normal;
    vec2 size;
    uint matIndex;
};


///
/// Base Tool
///
struct Ray{
    vec3 d;
    vec3 o;
};

vec3 ray(in Ray r,float t){
    return r.o+t*r.d;
}

/////////////////////////////////////
// the perspective camera          //
/////////////////////////////////////
struct ViewInfo{
    float f;// focal length of the camera
    mat4 mv;// view matrix
};

// Calculate fovy from focal length
float focalLengthToFovY(float f,float h){
    return 2.*atan(.5*h/f);
}

// Calculate focal length from fovy
float fovYToFocalLength(float fovy,float h){
    return h/(2.*tan(.5*fovy));
}

// Look at matrix
mat4 LookAt(in vec3 eye,in vec3 target,in vec3 up){
    vec3 f=normalize(eye-target);
    vec3 s=normalize(cross(up,f));
    vec3 u=normalize(cross(f,s));
    mat3 m=transpose(mat3(s,u,f));
    return mat4(vec4(m[0],0.),vec4(m[1],0.),vec4(m[2],0.),
        vec4(-m*eye,1.));    
}

ViewInfo CreateACamera(in vec3 e,in vec3 t,in float h,in float fovy){
    return ViewInfo(fovYToFocalLength(fovy,h),LookAt(e,t,L_F));
}

// Ray Casting
Ray CastRay(in vec2 p,in vec2 res,in float focalLength){
    return Ray(normalize(
            p.x*L_S*res.x+p.y*L_U*res.y+vec3(-res/2.,-focalLength)),
            vec3(0)
        );
}

/////////////////////////////////////
/////////////////////////////////////

///
/// Intersection
///
struct HitRecord{
    vec3 normal;
    vec3 pos;
    float t;
    vec3 color;    // Record current point surface color
    uint matIndex; // the material index
};

///
// Plane intersection test
// equation : t = dot((a-ro),n)/dot(rd,n) 
// note: a is the origin of the plane
///
bool IntersectWithPlane(in VS_ Ray r,in VS_ Plane p,out HitRecord hit){
    vec3 asubo=p.normal*p.offset-r.o;
    float d_dot_n=dot(r.d,p.normal);
    if(abs(d_dot_n)<EPS) return false;
    float hit_t=dot(asubo,p.normal)/d_dot_n;
    // Test if the t is in the range (TMIN,TMAX)
    bool hitted=false;
    if(hit_t>TMAX||hit_t<TMIN) return false;
    vec3 ray_hit=r.o+hit_t*r.d;
    if((hit.t>hit_t)){
        hit.normal=p.normal;
        hit.matIndex=p.matIndex;
        hit.pos=ray_hit;
        hit.t=hit_t;
        hitted=true;
    }
    return hitted;
}

bool IntersectWithRectangle(in VS_ Ray r,in VS_ Rectangle rec,out HitRecord hit){
    vec3 asubo=rec.origin-r.o;
    float d_dot_n=dot(r.d,rec.normal);
    if(abs(d_dot_n)<EPS) return false;
    float hit_t=dot(asubo,rec.normal)/d_dot_n;
    // Test if the t is in the range (TMIN,TMAX)
    bool hitted=false;
    if(hit_t>TMAX||hit_t<TMIN) return false;
    vec3 ray_hit=r.o+hit_t*r.d;
    if(hit.t>hit_t){
        vec3 y=rec.normal;
        vec3 x=normalize(cross(rec.normal,-L_F));
        vec3 z=normalize(cross(y,-x));
        mat3 basis=transpose(mat3(x,y,z));
        mat4 transform=mat4(vec4(basis[0],0.),vec4(basis[1],0.),vec4(basis[2],0.),
            vec4(-basis*rec.origin,1.));
        vec3 rhit=(transform*vec4(ray_hit,1.)).xyz;
        bvec2 re=lessThan(abs(rhit.xz),rec.size);
        if(re.x&&re.y){
            hit.normal=rec.normal;
            hit.matIndex=rec.matIndex;
            hit.pos=ray_hit;
            hit.t=hit_t;
            hitted=true;
        }
    }
    return hitted;
}

// Sphere intersection test
bool IntersectWithSphere(in VS_ Ray r,in VS_ Sphere s,out HitRecord hit)
{
    bool hitted=false;
    vec3 tmp=r.o-s.center;
    float b=2.*dot(r.d,tmp);
    float c=dot(tmp,tmp)-s.radius*s.radius;
    float discriminant=b*b-4.*c;
    if(discriminant>0.){
        discriminant=sqrt(discriminant);
        float t=(-b-discriminant)/(2.);
        if(t<TMIN) t=(-b+discriminant)/2.;
        if(t<TMIN||t>TMAX) return false;
        if(hit.t>t){
            hitted=true;
            hit.pos=r.o+r.d*t;
            hit.t=t;
            hit.normal=normalize(hit.pos-s.center);
            hit.matIndex=s.matIndex;
        }
        
    }
    return hitted;
}

///
/// Random utility
///
// From https://www.shadertoy.com/view/Xt3cDn
#if 1
//Modified from: iq's "Integer Hash - III" (https://www.shadertoy.com/view/4tXyWN)
//Faster than "full" xxHash and good quality
uint baseHash(uvec2 p)
{
    p = 1103515245U*((p >> 1U)^(p.yx));
    uint h32 = 1103515245U*((p.x)^(p.y>>3U));
    return h32^(h32 >> 16);
}

uint baseHash(uint p)
{
    p = 1103515245U*((p >> 1U)^(p));
    uint h32 = 1103515245U*((p)^(p>>3U));
    return h32^(h32 >> 16);
}
#else
//XXHash32 based (https://github.com/Cyan4973/xxHash)
//Slower, higher quality
uint baseHash(uvec2 p)
{
    const uint PRIME32_2 = 2246822519U, PRIME32_3 = 3266489917U;
	const uint PRIME32_4 = 668265263U, PRIME32_5 = 374761393U;
    uint h32 = p.y + PRIME32_5 + p.x*PRIME32_3;
    h32 = PRIME32_4*((h32 << 17) | (h32 >> (32 - 17))); //Initial testing suggests this line could be omitted for extra perf
    h32 = PRIME32_2*(h32^(h32 >> 15));
    h32 = PRIME32_3*(h32^(h32 >> 13));
    return h32^(h32 >> 16);
}

uint baseHash(uint p)
{
	const uint PRIME32_2 = 2246822519U, PRIME32_3 = 3266489917U;
	const uint PRIME32_4 = 668265263U, PRIME32_5 = 374761393U;
	uint h32 = p + PRIME32_5;
	h32 = PRIME32_4*((h32 << 17) | (h32 >> (32 - 17))); //Initial testing suggests this line could be omitted for extra perf
    h32 = PRIME32_2*(h32^(h32 >> 15));
    h32 = PRIME32_3*(h32^(h32 >> 13));
    return h32^(h32 >> 16);
}
#endif

float hash1(uint x)
{
    uint n = baseHash(x);
    return float((n >> 1) & (0x7fffffffU))/float(0x7fffffff);
}

vec2 hash22(uvec2 x)
{
    uint n = baseHash(x);
    uvec2 rz = uvec2(n, n*48271U);
    return vec2((rz.xy >> 1) & uvec2(0x7fffffffU))/float(0x7fffffff);
}

vec2 hash21(uint x)
{
    uint n = baseHash(x);
    uvec2 rz = uvec2(n, n*48271U); //see: http://random.mat.sbg.ac.at/results/karl/server/node4.html
    return vec2((rz.xy >> 1) & uvec2(0x7fffffffU))/float(0x7fffffff);
}

vec3 hash32(uvec2 x)
{
    uint n = baseHash(x);
    uvec3 rz = uvec3(n, n*16807U, n*48271U);
    return vec3((rz >> 1) & uvec3(0x7fffffffU))/float(0x7fffffff);
}

uint g_seed = 0u;

// frameIdx: When camera change this will change 
void InitGlobalSeed(uvec2 p, uint w, uint subFrameIdx) {
    g_seed = baseHash(uvec2(p.x*w+p.y, subFrameIdx));
}

///
/// Generate Random Ray Per Pixel
///

// Sample Per Pixel
#define SPP 4u

// Progressive processing
#define PROGRESSIVE 1
#define NON_UNIFORM
//#define NROOT 
//#define JITTLED
//#define MULTI_JITTLED


// Simple non uniform random sample
// Supper Sampling
void SimpleNonUniformDisRandom(in vec2 p,out vec2[SPP] samples){
    // random generate from 
    for(uint i=0u;i<SPP;i++){
        samples[i]=hash22((i+1u)*uvec2(p+p*samples[i]));
    }
}

void ProgressiveSNDR(out vec2 s){
    s=hash21(g_seed);
}

///
/// Stratified Random Sample
///

// N-root sample per pixel
void NRootSamplePP(in vec2 p,out vec2[SPP] samples){
    // First step: Generate random samples for a diagonal per pixel
    for(uint i=0u;i<SPP;i++){
        samples[i]=(vec2(i)+vec2(-.5)+hash22((i+1u)*uvec2(p+p*samples[i])))/vec2(SPP);
    }
    // Second step: Shuffle The random samples
    for(uint s=SPP-2u;s>=0u;s--){
        // Generate a random index for current sample;
        uint k=uint(baseHash((s*50054430u))%uint(SPP));
        float tmp=samples[s+1u].x;
        samples[s+1u].x=samples[k].x;
        samples[k].x=tmp;
    }
}

// Jittled Sample per pixel
void JittledSamplePP(in vec2 p,out vec2[SPP] samples){
    uint sqrt_samples=uint(sqrt(float(SPP)));
    for(uint i=0u;i<sqrt_samples;i++)
        for(uint j=0u;j<sqrt_samples;j++){
            vec2 rand=vec2(-.5)+hash22((i+1u)*uvec2(p));
            samples[i*sqrt_samples+j]=(vec2(i,j)+rand)/vec2(sqrt_samples);
        }
}

// Multi-Jittled Sample per pixel
// combined n-roots with jittled sample techniques
void MutltiJittledSamplePP(in vec2 p,out vec2[SPP] samples){
    uint sqrt_samples=uint(sqrt(float(SPP)));
    float sub_cell_width=1./float(SPP);
    for(uint i=0u;i<sqrt_samples;i++)
        for(uint j=0u;j<sqrt_samples;j++){
            vec2 rand=vec2(-.5)+hash22(uvec2(i+1u)*uvec2(p));
            samples[i*sqrt_samples+j]=vec2(
                float(i)*sub_cell_width*float(sqrt_samples)+float(j)*sub_cell_width+rand.x*sub_cell_width,
                float(j)*sub_cell_width*float(sqrt_samples)+float(i)*sub_cell_width+rand.y*sub_cell_width
            );
        }
    // shuffle the xy coordinates
    for(uint i=0u;i<sqrt_samples;i++)
        for(uint j=0u;j<sqrt_samples;j++){
            uvec2 k=uvec2(j)+uvec2(
                baseHash(j*50054430u)%(1u+uint(sqrt_samples-j-1u)),
                baseHash(j*30352430u)%(1u+uint(sqrt_samples-j-1u))
            );
            float t=samples[i*sqrt_samples+j].x;
            samples[i*sqrt_samples+j].x=samples[i*sqrt_samples+k.x].x;
            samples[i*sqrt_samples+k.x].x=t;

            t=samples[j*sqrt_samples+i].y;
            samples[j*sqrt_samples+i].y=samples[k.y*sqrt_samples+i].y;
            samples[k.y*sqrt_samples+i].y=t;
        }
}


///
/// Sample Method
///

//#define SIMPLE_UNIFORM_ 
#define COSINE_UNIFORM_ 

vec2 DiskUnitSample(float seed) {
    vec2 rand=hash22(uvec2(seed))*vec2(1.,2.*PI);
    float r=rand.x;
    float phi=rand.y;
    return vec2(r*cos(phi),r*sin(phi));
}

vec3 hemisphereSample(float seed){
    vec3 sampleVec=vec3(0.);
    // reference: https://mathworld.wolfram.com/SpherePointPicking.html
#if defined(SIMPLE_UNIFORM_)
    vec2 rand=hash21(g_seed)*vec2(PI,2.*PI);
    vec2 theta=vec2(sin(rand.x),cos(rand.x));
    vec2 phi=vec2(sin(rand.y),cos(rand.y));
    sampleVec=vec3(theta.x*phi.y,theta.y,theta.x*phi.x);
#elif defined(COSINE_UNIFORM_)
    vec2 rand=hash21(g_seed)*vec2(PI,2.*PI);
    float r = sqrt(rand.x);
    float theta = 2.*PI*rand.y;
 
    float x = r*cos(theta);
    float y = r*sin(theta);
    sampleVec=vec3(x,y,sqrt(max(0.f,1.-rand.x*rand.x-rand.y*rand.y)));
#else

#endif
    return sampleVec;
}

//////////////////////////
// single direction path tracing and ray tracing
//////////////////////////

///
// Material
///



#define LAMBERTIAN   0u
#define COOKTORRANCE 1u
#define ORENLAYER    2u
#define PHONG        3u

struct Material{
    uint type;
    vec3 baseColor;
    /// diffuse relectance
    vec3 kd;
    
    vec3 ks;
};

// Radiance Definetion
struct Radiance {
    vec3 emitted;      // emitted radiance 
    vec3 attenuation;  // note: fr term
    vec3 radiance;     // radiance existance
};

#define DEFAULT_MAT 0u

// Define material
const uint materialCnt=5u;
Material materials[]=Material[materialCnt](
    Material(LAMBERTIAN,vec3(0.),vec3(1.),vec3(0.)), // Default material
    Material(LAMBERTIAN,vec3(0.),vec3(0.,1.,0.),vec3(0.)),
    Material(LAMBERTIAN,vec3(0.),vec3(0.7,0.32,0.22),vec3(0.)),
    Material(LAMBERTIAN,vec3(0.),vec3(1.,0.,0.),vec3(0.)),
    Material(PHONG,vec3(0.),vec3(0.2,0.662,0.92),vec3(0.3))
);

/// 
// Light Setup
///
#define POINT_L 0

#define LT_DIRECT 0u  // Default light type
#define LT_POINT  1u

struct Light{
    uint type;
    vec3 color;      // the rgb intensity of the light
    float intensity; // the hole luminance 
    vec3 dir;        // direction of the light   
#if POINT_L 
    vec3 pos; // the position of the point light source
#endif
};

const uint lightCnt=1u;
Light lights[]=Light[lightCnt](
    // A Directional light
    Light(
        LT_DIRECT,vec3(1.),10.,vec3(normalize(vec3(0.)-vec3(2.,-1.,1.)))
#if POINT_L
        ,vec3(vec3(1.))
#endif
    )
);

///
// Radiance Setup
///
struct IRRadiance{
    vec3 irradiance;
    vec3 incident;
};

struct BRDFdata{
    vec3 frd;
    vec3 frs;
};

// Bidirectional Scattering Distribution Function
vec3 BSDF(in HitRecord h){
    Material mat=materials[h.matIndex];
#if defined(LAMBERTIAN)
    /// lambertian brdf
    //  fr(x,d)=kd/PI
    ///
    if(mat.type==LAMBERTIAN) 
        return mat.kd/PI;
#endif

#if defined(PHONG)
    if(mat.type==PHONG) //TODO: Implement
        
#endif

#if defined(COOKTORRANCE)
    if(mat.type==COOKTORRANCE) //TODO: Implement
#endif

#if defined(ORENLAYER)
    if(mat.type==ORENLAYER) //TODO: Implement
#endif

    return vec3(0.);
}

// direct illumination and indirect illumination
vec3 L_existance_over_hemisphere(in Light li,in vec3 col,in vec3 d,in float v,out HitRecord h){
#if GI

#else
    // L direct
    h.color=BSDF(h)*col;
    //h.color=BSDF(h);
    //h.color*=v;
#endif
    return h.color;
}

/// incident radiance,integeration over hemisphere
vec3 L_irradiance_over_hemisphere(in Light li,in vec3 d,out HitRecord h){
    // Setting 
#if GI

#else
    // Simple direct illumination 
    // note: not include other terms like ambient term and specular term
    //h.color*=materials[h.matIndex].kd;
    // 
    h.color*=BSDF(h)*max(0.,dot(-d,h.normal))+BSDF(h)*li.color*li.intensity*max(0.,dot(-li.dir,h.normal));
#endif
    return h.color;
}

/// incident radiance,integeration over surfaces
vec3 L_irradiance_over_surface(in Light li,in HitRecord h){
    vec3 o=vec3(0.);
    // Setting 
#if GI

#else
    
#endif
    return o;
}
